#pragma once

#include "stdafx.h"
#include <list>
#include <Windows.h>
#include "PEFile.h"

typedef std::list<IMAGE_SECTION_HEADER> SECTION_LIST;

class CEnumPECallback;

class CPEParse : public CPEFile
{
public:
    CPEParse(void);
    virtual ~CPEParse(void);

public:
    BOOL 
    Initialize(
        VOID
    );

    VOID 
    Free(
        VOID
    );

    PIMAGE_DOS_HEADER 
    GetDosHeader(
        VOID
    );

    PIMAGE_NT_HEADERS 
    GetNtHeader(
        VOID
    );

    BOOL 
    IsRvaInSection(
        IN ULONG Rva, 
        IN PIMAGE_SECTION_HEADER SectionHead
    );

    BOOL 
    GetTimeStamp(
        OUT PULONG ts
    );

    BOOL 
    IsVaInSection(
        IN ULONG Va, 
        IN PIMAGE_SECTION_HEADER SectionHead
    );

    ULONG 
    RvatoVa(
        IN ULONG Rva
    );

    ULONG 
    VatoRva(
        IN ULONG Va
    );

    ULONG 
    GetRealAddress(
        IN ULONG Address
    );

    ULONG 
    DirectoryEntryToData(
        IN USHORT Index, 
        OUT PULONG Size
    );

    ULONG 
    GetBaseAddress(
        VOID
    );

    int 
    PEstrcmp(
        IN ULONG Offset, 
        IN LPCSTR Src
    );

    int 
    PEwcscmp(
        IN ULONG Offset, 
        IN LPCWSTR Src
    );

    int 
    PEwcsnnmp(
        IN ULONG Offset, 
        IN LPCWSTR Src, 
        IN ULONG Count
    );

    ULONG 
    LookupExportRoutineById(
        IN ULONG Id
    );

    ULONG 
    LookupExportRoutineByName(
        IN LPCSTR lpszName
    );

    PVOID 
    RelocationSection(
        IN ULONG RelocatBase, 
        IN PIMAGE_SECTION_HEADER Section, 
        OUT PULONG Length
    );
    
    VOID 
    SetEnumCallBack(
        IN CEnumPECallback* clsCallBack
    );

    BOOL 
    EnumRelocationData(
        VOID
    );

    ULONG
    GetFirstSectionOff(
        VOID
    );

    ULONG
    GetSectionCounts(
        VOID
    );

    ULONG
    PEstrlen(
        IN ULONGLONG Offset
    );

    ULONG
    PEwcslen(
        IN ULONGLONG Offset
    );

    LPSTR
    PeReadString(
        IN ULONGLONG Offset
    );

    static BOOL
    IsExportById(
        IN LPCSTR lpszFuncName
    );

    BOOL
    EnumExportFunction(
        VOID
    );

    ULONG
    LookupImportedRoutineIATRvaByName(
        IN LPCSTR lpszFuncName
    );

    ULONG
    LookupImportedRoutineIATRvaByName(
        IN LPCSTR lpszImgName,
        IN LPCSTR lpszFuncName
    );

    ULONG
    GetEntryPoint(
        VOID
    );

    SECTION_LIST&
    GetSectionList(
        VOID
    )
    {
        return m_Sections;
    }

    BOOL
    IsAddressInDirectoryRange(
        IN ULONG Address
    );

private:

    int
    IsIdInOrdinalTable(
        IN ULONG Id,
        IN PUSHORT OrdinalTable,
        IN ULONG Counts
    );

private:
    IMAGE_DOS_HEADER m_DosHeader;
    IMAGE_NT_HEADERS m_NtHeader;
    SECTION_LIST m_Sections;
    CEnumPECallback* m_clsEnumCallback;
};

class CEnumPECallback
{
public:
    CEnumPECallback(void){};
    virtual ~CEnumPECallback(void){};
    virtual BOOL EnumExportFunctionCallBack(
            IN CPEParse* clsPeFile,
            IN LPCSTR lpszFuncName,
            IN ULONG Address,
            IN ULONG FunctionId,
            IN LPCSTR lpszForward) = 0;
    virtual BOOL EnumImportFunctionCallBack() = 0;
    virtual BOOL EnumRelocationDataCallBack(
            IN CPEParse* clsPeFile,
            IN PIMAGE_BASE_RELOCATION pRelocBlock,
            IN ULONG RelocAddress
        ) = 0;
};